from dataclasses import dataclass
from collections import ChainMap


def get_day_name(n:int):
    if n<30:
        return "d" + str(n)
    if n<360:
        return "m" + str(n//30)
    if n < 10800:
        return "y" + str(n//360)
    return "tall"

@dataclass
class TimeWindow:
    
    window_days:list = None
    window_names:list = None
    
    def __post_init__(self):
        if self.window_days is None:
            self.window_days = [3, 7, 15, 30, 60, 90, 180, 50000]
            
            
        if self.window_names is None:
            self.window_names = [get_day_name(n) for n in self.window_days]
    
    def iter(self, data):
        for name, days in zip(self.window_names, self.window_days):
            yield name, data[(data['up2in_time_diff_days'] > 0) & (data['up2in_time_diff_days'] <= days)]
    
    def add_tags(self, up2in_time_diff_days: int):
        tags = []
        for name, days in zip(self.window_names, self.window_days):
            if up2in_time_diff_days < days:
                tags.append(name)
        return tags
    
    def create_features(self, data, func, *args, **kwargs):
        dfs = []
        for name, df in self.iter(data):
            features = func(df, *args, **kwargs)
            features.columns = [f"{col}_{name}" for col in features.columns]
            dfs.append(features)
        return self.merge_dfs(dfs)
    
    @staticmethod
    def merge_dfs(dfs):
        return dict(ChainMap(*dfs))
    
    @staticmethod
    def get_cycdata(data, cycrate):
        data1 = data[(data['up2in_time_diff_days'] > 0) & (data['up2in_time_diff_days'] <= cycrate)]
        data2 = data[(data['up2in_time_diff_days'] > cycrate) & (data['up2in_time_diff_days'] <= cycrate*2)]
        return data1, data2
    
if __name__ == '__main__':
    demo = TimeWindow([15, 30, 60, 90, 180, 50000])
    print(demo.window_days)
    print(demo.window_names)
    print(get_day_name(30))